from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import SGD
import matplotlib.pyplot as plt
import numpy as np
from skimage.transform import resize
from data_loader import load_data

# 这两句解决OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized.报错问题
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


(x_train, y_train), (x_test, y_test) = load_data()

print(x_train.shape)
print(y_train.shape)
print(x_test.shape)
print(y_test.shape)

x_train = x_train.reshape(x_train.shape[0], -1)
x_test = x_test.reshape(x_test.shape[0], -1)
x_train.shape

cat_model = Sequential()
cat_model.add(Dense(128, activation='relu', input_shape=(12288,)))
cat_model.add(Dense(64, activation='relu'))
cat_model.add(Dense(32, activation='relu'))
cat_model.add(Dense(16, activation='relu'))
cat_model.add(Dense(1, activation='sigmoid'))

cat_model.summary()

cat_model.compile(optimizer=SGD(), loss='binary_crossentropy', metrics=['accuracy'])

cat_model.fit(x_train, y_train, epochs=40, validation_data=(x_test, y_test))

fig = plt.figure(figsize=(16, 16))
for i in range(1, 17):
    my_image = 'images/test/{}.jpg'.format(i)
    my_image = np.array(plt.imread(my_image))
    ax = fig.add_subplot(4, 5, i)
    plt.imshow(my_image)
    num_px = 64
    my_image = resize(my_image, (num_px, num_px))
    my_image.shape
    my_image = my_image.reshape(1, -1)
    a = cat_model.predict(my_image)
    if a > 0.5:
        ax.title.set_text('cat {}'.format(a))
    else:
        ax.title.set_text('dog {}'.format(1 - a))

plt.show()
cat_model.evaluate(x_test, y_test)
